# @Time : 2021/8/5 9:52
# @Author : Li Kunlun
# @Description : 二维卷积层


from mxnet import autograd, nd
from mxnet.gluon import nn


# 1、二维互相关运算
def corr2d(X, K):
    h, w = K.shape
    Y = nd.zeros((X.shape[0] - h + 1, X.shape[1] - w + 1))
    for i in range(Y.shape[0]):
        for j in range(Y.shape[1]):
            Y[i, j] = (X[i: i + h, j: j + w] * K).sum()
    return Y


X = nd.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
K = nd.array([[0, 1], [2, 3]])
# [[ 19.  25.]
#  [ 37.  43.]]
# <NDArray 2x2 @cpu(0)>
print(corr2d(X, K))


# 2、二维卷积层
# 基于corr2d函数来实现一个自定义的二维卷积层。在构造函数__init__里我们声明weight和bias这两个模型参数
class Conv2D(nn.Block):
    def __init__(self, kernel_size, **kwargs):
        super(Conv2D, self).__init__(**kwargs)
        self.weight = self.params.get('weight', shape=kernel_size)
        self.bias = self.params.get('bias', shape=(1,))

    def forward(self, x):
        return corr2d(x, self.weight.data()) + self.bias.data()


# 3、图像中物体边缘检测
"""
1、它中间4列为黑（0），其余为白（1）
    [[ 1.  1.  0.  0.  0.  0.  1.  1.]
     [ 1.  1.  0.  0.  0.  0.  1.  1.]
     [ 1.  1.  0.  0.  0.  0.  1.  1.]
     [ 1.  1.  0.  0.  0.  0.  1.  1.]
     [ 1.  1.  0.  0.  0.  0.  1.  1.]
     [ 1.  1.  0.  0.  0.  0.  1.  1.]]
    <NDArray 6x8 @cpu(0)>

从白到黑的边缘和从黑到白的边缘分别检测成了1和-1:
    [[ 0.  1.  0.  0.  0. -1.  0.]
     [ 0.  1.  0.  0.  0. -1.  0.]
     [ 0.  1.  0.  0.  0. -1.  0.]
     [ 0.  1.  0.  0.  0. -1.  0.]
     [ 0.  1.  0.  0.  0. -1.  0.]
     [ 0.  1.  0.  0.  0. -1.  0.]]
    <NDArray 6x7 @cpu(0)>
"""
X = nd.ones((6, 8))
X[:, 2:6] = 0
print(X)

# 构造一个高和宽分别为1和2的卷积核K
K = nd.array([[1, -1]])
Y = corr2d(X, K)
print(Y)

# 4、通过数据学习核数组
# 构造一个输出通道数为1，核数组形状是(1, 2)的二维卷积层
conv2d = nn.Conv2D(1, kernel_size=(1, 2))
conv2d.initialize()

# 二维卷积层使用4维输入输出，格式为(样本, 通道, 高, 宽)，这里批量大小（批量中的样本数）和通道数均为1
X = X.reshape((1, 1, 6, 8))
Y = Y.reshape((1, 1, 6, 7))

for i in range(10):
    with autograd.record():
        Y_hat = conv2d(X)
        l = (Y_hat - Y) ** 2
    l.backward()
    # 简单起见，这里忽略了偏差
    conv2d.weight.data()[:] -= 3e-2 * conv2d.weight.grad()
    if (i + 1) % 2 == 0:
        """ 10次迭代后误差已经降到了一个比较小的值:
        batch 2, loss 4.949
        batch 4, loss 0.831
        batch 6, loss 0.140
        batch 8, loss 0.024
        batch 10, loss 0.004
        """
        print('batch %d, loss %.3f' % (i + 1, l.sum().asscalar()))

"""
1、学习到的K核数组输出
    [[ 0.98949999 -0.98737049]]
    <NDArray 1x2 @cpu(0)>
"""
print(conv2d.weight.data().reshape((1, 2)))
